package com.aigc.config;

import cn.hutool.core.collection.CollUtil;
import com.aigc.domin.Conversation;
import com.aigc.mapper.ConversationMapper;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.time.LocalDateTime;
import java.util.List;

@Component
public class MyChatMemory implements ChatMemory {

    @Autowired
    private ConversationMapper conversationMapper;

    @Override
    public void add(String conversationId, List<Message> messages) {
        messages.forEach(m -> {
            Conversation conversation = Conversation.builder()
                    .conversationId(conversationId)
                    .message(m.getText())
                    .type(m.getMessageType().getValue())
                    .createTime(LocalDateTime.now())
                    .build();
            conversationMapper.insert(conversation);
        });
    }

    @Override
    public List<Message> get(String conversationId) {
        List<Conversation> conversations = conversationMapper.selectList(new LambdaQueryWrapper<Conversation>()
                .eq(Conversation::getConversationId, conversationId)
        );

            return conversations.stream().map(c -> {
                Message message = null;
                if (MessageType.USER.getValue().equals(c.getType())) {
                    message = new UserMessage(c.getMessage());
                } else if (MessageType.ASSISTANT.getValue().equals(c.getType())) {
                    message = new AssistantMessage(c.getMessage());
                }
                return message;
            }).toList();

    }

    @Override
    public void clear(String conversationId) {
        conversationMapper.delete(new LambdaQueryWrapper<Conversation>()
                .eq(Conversation::getConversationId, conversationId));
    }
}
